#include <iostream>
#include <fstream>
#include <cstdlib>
#include <string>
#include <cstdio>
#include "antlr4-runtime/antlr4-runtime.h"
#include "DecafLangLexer.h"
#include "DecafLangParser.h"
#include "ASTVisitor.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/raw_ostream.h"

using namespace decaf_lang::parser;
using namespace antlr4;

int main(int argc, const char* argv[]) {
    if (argc < 3) {
        std::cerr << "Usage: " << argv[0] << " [build|run] <input_file.decaf> [QUIET]" << std::endl;
        return 1;
    }

    std::string mode = argv[1];
    std::string inputFile = argv[2];
    bool quiet = (argc > 3 && std::string(argv[3]) == "QUIET");

    std::string base = inputFile.substr(0, inputFile.find_last_of('.'));
    std::string llFile = base + ".ll";
    std::string binFile = base + ".bin";

    std::ifstream stream(inputFile);
    if (!stream) {
        if (!quiet) std::cerr << "Error: Could not open file " << inputFile << std::endl;
        return 1;
    }

    ANTLRInputStream input(stream);
    DecafLangLexer lexer(&input);
    CommonTokenStream tokens(&lexer);
    DecafLangParser parser(&tokens);

    tree::ParseTree* tree = parser.program();

    llvm::LLVMContext context;
    llvm::Module module("DecafLangModule", context);
    llvm::IRBuilder<> builder(context);

    ASTVisitor visitor(context, module, builder);
    try {
        visitor.visit(tree);
    } catch (const std::exception& e) {
        if (!quiet) std::cerr << "Compilation Error: " << e.what() << std::endl;
        return 1;
    }

    // Write IR to .ll file
    std::error_code EC;
    llvm::raw_fd_ostream outFile(llFile, EC);
    module.print(outFile, nullptr);
    outFile.close();
    if (!quiet) std::cout << "Wrote LLVM IR to " << llFile << std::endl;

    // Compile to native binary with -O3
    std::string clang = "/opt/homebrew/opt/llvm@17/bin/clang";
    std::string cmd = clang + " -O3 " + llFile + " -o " + binFile;
    if (quiet) cmd += " >/dev/null 2>&1";
    if (!quiet) std::cout << "Optimizing and compiling: " << cmd << std::endl;
    int ret = std::system(cmd.c_str());
    if (ret != 0) {
        if (!quiet) std::cerr << "clang failed!" << std::endl;
        return 1;
    }
    if (!quiet) std::cout << "Wrote optimized native binary to " << binFile << std::endl;

    if (mode == "run") {
        if (!quiet) std::cout << "Running: " << binFile << std::endl;
        ret = std::system(binFile.c_str());
        return ret;
    }

    return 0;
}
